"""离群值处理"""

import pandas as pd
import numpy as np
from tools.base_tools import conv_arr_to_df




def winsorize_outliers_with_mad(df, columns=None, n=3, caches=None):
    """
    使用MAD方法对指定列的离群值进行上下限调整 (Winsorization)。

    参数:
    - df: pandas.DataFrame，输入的数据框
    - columns: list，指定需要处理的列名
    - n: float，离群值的倍数阈值，默认是3

    返回:
    - pandas.DataFrame，处理后的数据框
    """
    df = conv_arr_to_df(df, columns)
    df_1 = df.copy()  # 防止修改原始数据

    caches = caches if caches else {}
    for col in columns:
        if col not in df_1.columns:
            raise ValueError(f"列 '{col}' 不在 DataFrame 中")

        cache_i = caches.get(col)
        if cache_i:
            lower_bound, upper_bound = cache_i.get('lower_bound'), cache_i.get('upper_bound')
        else:
            # 计算中位数和MAD
            median = df_1[col].median()
            absolute_deviation = abs(df_1[col] - median)
            mad = absolute_deviation.median()

            # 离群值上下限
            lower_bound = median - n * mad
            upper_bound = median + n * mad

            cache_i = {'lower_bound': lower_bound, 'upper_bound': upper_bound}
            caches[col] = cache_i

        # 调整离群值至上下限
        df_1[col] = df_1[col].apply(lambda x: max(lower_bound, min(upper_bound, x)))

        outliers_nums = sum(df_1[col] != df[col])
        print(f'--- col[{col}] --- lower_bound, upper_bound: [{lower_bound: .4f}, {upper_bound: .4f}] --- outliers_nums: {outliers_nums} / {df.shape[0]}')

    return df_1, caches


def winsorize_outliers_with_sigma(df, columns=None, n=3, caches=None):
    """
    使用 \(\sigma\) 方法对指定列的离群值进行上下限调整 (Winsorization)。

    参数:
    - df: pandas.DataFrame，输入的数据框
    - columns: list，指定需要处理的列名
    - n: float，离群值的倍数阈值，默认是3

    返回:
    - pandas.DataFrame，处理后的数据框
    """
    if not isinstance(df, pd.DataFrame):
        df = pd.DataFrame(df, columns=columns)
    df_1 = df.copy()  # 防止修改原始数据
    caches = caches if caches else {}
    for col in columns:
        if col not in df_1.columns:
            raise ValueError(f"列 '{col}' 不在 DataFrame 中")

        cache_i = caches.get(col)
        if cache_i:
            lower_bound, upper_bound = cache_i.get('lower_bound'), cache_i.get('upper_bound')
        else:
            # 计算均值和标准差
            mean = df_1[col].mean()
            std_dev = df_1[col].std()

            # 离群值上下限
            lower_bound = mean - n * std_dev
            upper_bound = mean + n * std_dev

            cache_i = {'lower_bound': lower_bound, 'upper_bound': upper_bound}
            caches[col] = cache_i

        # 调整离群值至上下限
        df_1[col] = df_1[col].apply(lambda x: max(lower_bound, min(upper_bound, x)))
        # df_1[col].describe()

        outliers_nums = sum(df_1[col] != df[col])
        print(f'--- col[{col}] --- lower_bound, upper_bound: [{lower_bound: .4f}, {upper_bound: .4f}] --- outliers_nums: {outliers_nums}  / {df.shape[0]}')

    return df_1, caches


if __name__ == '__main__':
    # 示例数据
    data = {
        "A": [10, 12, 13, 14, 15, 100],
        "B": [1, 2, 2, 3, 1, 50],
        "C": [5, 7, 8, 6, 7, 6]
    }
    df = pd.DataFrame(data)

    arr = np.array(df)
    columns = None
    processed_df = winsorize_outliers_with_mad(arr, columns=['a', 'b', 'c'], n=3)

    # 对列 A 和 B 进行处理
    # processed_df = winsorize_outliers_with_mad(df, columns=["A", "B", "C"], n=3)
    # processed_df = winsorize_outliers_with_sigma(df, columns=["A", "B", "C"], n=1)

    print("处理前的数据：\n", df)
    print("处理后的数据：\n", processed_df)

